from Crypto.Random import get_random_bytes
from Crypto.Cipher import AES
import base64, json

class AESCipher:
    def __init__(self, key: bytes = None):
        # 256-bit 密钥 = 32 bytes
        self.key = key or get_random_bytes(32)

    def encrypt(self, plaintext: str) -> str:
        """返回 base64 字符串，包含 nonce + ciphertext + tag"""
        cipher = AES.new(self.key, AES.MODE_GCM)
        ct_bytes, tag = cipher.encrypt_and_digest(plaintext.encode())
        blob = json.dumps({
            'n': base64.b64encode(cipher.nonce).decode(),
            't': base64.b64encode(tag).decode(),
            'c': base64.b64encode(ct_bytes).decode()
        })
        return base64.b64encode(blob.encode()).decode()

    def decrypt(self, b64data: str) -> str:
        """解密失败会抛出 ValueError（认证失败）"""
        blob = json.loads(base64.b64decode(b64data))
        nonce = base64.b64decode(blob['n'])
        tag   = base64.b64decode(blob['t'])
        ct    = base64.b64decode(blob['c'])
        cipher = AES.new(self.key, AES.MODE_GCM, nonce=nonce)
        return cipher.decrypt_and_verify(ct, tag).decode()
    @staticmethod
    def b64_encode(data: str | bytes, encoding='utf-8') -> str:
        """
        任意字符串 → base64 字符串
        """
        if isinstance(data, str):
            data = data.encode(encoding)
        return base64.b64encode(data).decode('ascii')

    @staticmethod
    def b64_decode(b64_str: str, encoding='utf-8') -> str:
        """
        base64 字符串 → 原文
        """
        raw_bytes = base64.b64decode(b64_str.encode('ascii'))
        return raw_bytes.decode(encoding)


import base64, json, os
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Protocol.KDF import PBKDF2

SALT_SIZE = 16          # 128-bit salt
KEY_SIZE  = 32          # AES-256
PBKDF2_ROUNDS = 100_000  # 可调，越高越慢

class SaltedB64Cipher:
    @staticmethod
    def encrypt(plaintext: str, password: str) -> str:
        salt = get_random_bytes(SALT_SIZE)
        key = PBKDF2(password.encode(), salt, dkLen=KEY_SIZE, count=PBKDF2_ROUNDS)

        cipher = AES.new(key, AES.MODE_GCM)
        ct_bytes, tag = cipher.encrypt_and_digest(plaintext.encode())

        # 打包成字典 -> json -> base64
        blob = {
            'salt': base64.b64encode(salt).decode(),
            'nonce': base64.b64encode(cipher.nonce).decode(),
            'tag': base64.b64encode(tag).decode(),
            'ct': base64.b64encode(ct_bytes).decode()
        }
        return base64.b64encode(json.dumps(blob).encode()).decode()

    @staticmethod
    def decrypt(b64_str: str, password: str) -> str:
        blob = json.loads(base64.b64decode(b64_str))
        salt  = base64.b64decode(blob['salt'])
        nonce = base64.b64decode(blob['nonce'])
        tag   = base64.b64decode(blob['tag'])
        ct    = base64.b64decode(blob['ct'])

        key = PBKDF2(password.encode(), salt, dkLen=KEY_SIZE, count=PBKDF2_ROUNDS)
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
        return cipher.decrypt_and_verify(ct, tag).decode()

# 用法
